from ultralytics import YOLO
import cv2
import os
import numpy as np
import time
from datetime import datetime
import pandas as pd
from openpyxl import Workbook


class HighwayVehicleDetector:
    def __init__(self, model_path='yolov8n.pt'):
        """初始化检测器
        model_path: 模型权重路径，默认使用YOLOv8n预训练模型（轻量级，适合快速部署）
        """
        # 加载YOLOv8模型（预训练模型已在COCO数据集上训练，包含80类目标）
        self.model = YOLO(model_path)

        # 定义类别映射（COCO数据集类别）
        self.class_names = self.model.names  # 模型自带的类别名称字典

        # 定义"汽车"类别（COCO数据集中汽车对应ID=2）
        self.car_class_id = 2

        # 非汽车类别提醒话术（根据常见高速异常目标定义）
        self.warning_messages = {
            0: "行人：高速禁止行人进入，存在严重安全隐患！",  # 行人
            1: "自行车：非机动车禁止上高速，请及时处理！",  # 自行车
            3: "摩托车：部分高速禁止摩托车通行，请注意核查！",  # 摩托车
            4: "公交车：属于汽车范畴（计数）",  # 公交车（补充：虽然是汽车，但可根据需求调整）
            # 其他可能的非汽车类别（如动物、物品等）
            16: "猫：动物闯入高速，需警惕碰撞风险！",  # 猫
            17: "狗：动物闯入高速，需警惕碰撞风险！",  # 狗
            72: "瓶子：路面有散落物品，可能影响行车安全！",  # 瓶子
            73: "杯子：路面有散落物品，可能影响行车安全！",  # 杯子
            "default": "非汽车目标：检测到未定义的非汽车物体，请注意！"
        }

    def detect_and_process(self, image_path, conf_threshold=0.5, save_result=True, result_filename=None):
        """检测图片中的目标并处理结果
        image_path: 图片路径
        conf_threshold: 置信度阈值，过滤低置信度检测结果
        save_result: 是否保存检测结果图片
        result_filename: 保存的结果文件名，不指定则自动生成
        return: 汽车计数、提醒信息列表、处理时间
        """
        start_time = time.time()  # 记录开始时间
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片：{image_path}")

        # 模型推理（返回检测结果）
        results = self.model(image, conf=conf_threshold)

        car_count = 0  # 汽车计数器
        warnings = []  # 非汽车提醒列表

        # 解析检测结果
        for result in results:
            # 遍历每个检测到的目标
            for box in result.boxes:
                class_id = int(box.cls[0])  # 目标类别ID
                class_name = self.class_names[class_id]  # 目标类别名称
                conf = float(box.conf[0])  # 置信度

                # 判断是否为汽车
                if class_id == self.car_class_id:
                    car_count += 1
                    print(f"检测到汽车（置信度：{conf:.2f}）")
                else:
                    # 非汽车，生成提醒信息
                    message = self.warning_messages.get(class_id, self.warning_messages["default"])
                    warnings.append(f"[{class_name}] {message}（置信度：{conf:.2f}）")
                    print(f"检测到非汽车：{class_name}（置信度：{conf:.2f}）")

        # 绘制检测结果（可选，用于可视化）
        if save_result:
            annotated_image = results[0].plot()  # 自动标注边界框和类别
            output_dir = "imgs/resultImage"
            
            # 确保输出目录存在
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            # 生成保存文件名
            if result_filename is None:
                # 获取原始文件名（不含扩展名）
                original_filename = os.path.splitext(os.path.basename(image_path))[0]
                result_filename = f"{original_filename}_detected.jpg"
            
            # 拼接完整保存路径
            output_path = os.path.join(output_dir, result_filename)
            
            # 保存到指定文件夹
            cv2.imwrite(output_path, annotated_image)
            print(f"标注结果已保存为：{output_path}")
        
        # 计算处理时间
        processing_time = time.time() - start_time
        
        return car_count, warnings, processing_time
    
    def batch_process_images(self, image_folder, conf_threshold=0.5, save_results=True):
        """批量处理文件夹中的所有图片并计算车流量
        image_folder: 包含图片的文件夹路径
        conf_threshold: 置信度阈值
        save_results: 是否保存每张图片的检测结果
        return: 处理结果统计信息
        """
        # 获取文件夹中的所有图片文件
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
        image_files = [f for f in os.listdir(image_folder) 
                      if os.path.isfile(os.path.join(image_folder, f)) and 
                      os.path.splitext(f)[1].lower() in image_extensions]
        
        if not image_files:
            print(f"在文件夹 {image_folder} 中未找到图片文件")
            return None
        
        print(f"找到 {len(image_files)} 张图片待处理...")
        
        # 初始化统计信息
        total_cars = 0
        total_images = len(image_files)
        total_processing_time = 0
        image_results = []
        
        # 处理每张图片
        for i, image_file in enumerate(image_files, 1):
            image_path = os.path.join(image_folder, image_file)
            print(f"\n处理图片 {i}/{total_images}: {image_file}")
            
            try:
                # 检测处理图片
                car_count, warnings, processing_time = self.detect_and_process(
                    image_path, conf_threshold, save_results
                )
                
                # 更新统计信息
                total_cars += car_count
                total_processing_time += processing_time
                
                # 保存单张图片的结果
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                image_results.append({
                    'filename': image_file,
                    'timestamp': timestamp,
                    'car_count': car_count,
                    'processing_time': processing_time,
                    'warnings': warnings
                })
                
                print(f"图片 {image_file} 处理完成，检测到 {car_count} 辆汽车")
                
            except Exception as e:
                print(f"处理图片 {image_file} 时出错：{str(e)}")
        
        # 计算平均车流量和平均处理时间
        avg_cars_per_image = total_cars / total_images if total_images > 0 else 0
        avg_processing_time = total_processing_time / total_images if total_images > 0 else 0
        
        # 生成统计报告
        stats = {
            'total_images': total_images,
            'total_cars': total_cars,
            'avg_cars_per_image': avg_cars_per_image,
            'total_processing_time': total_processing_time,
            'avg_processing_time': avg_processing_time,
            'image_results': image_results
        }
        
        # 打印统计报告
        self._print_traffic_stats(stats)
        
        # 保存统计结果到文件（包括Excel）
        self._save_traffic_stats(stats)
        
        return stats
    
    def _print_traffic_stats(self, stats):
        """打印车流量统计信息"""
        print("\n" + "="*50)
        print("车流量统计报告")
        print("="*50)
        print(f"处理图片总数: {stats['total_images']}")
        print(f"检测到汽车总数: {stats['total_cars']}")
        print(f"平均每图片汽车数: {stats['avg_cars_per_image']:.2f}")
        print(f"总处理时间: {stats['total_processing_time']:.2f} 秒")
        print(f"平均处理时间: {stats['avg_processing_time']:.2f} 秒/图片")
        print("="*50)
    
    def _save_traffic_stats(self, stats):
        """保存车流量统计信息到文件"""
        # 生成时间戳用于文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_dir = "logs"
        
        # 确保日志目录存在
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
        
        # 创建数据列表用于Excel写入
        excel_data = []
        
        # 保存详细统计到CSV文件
        csv_file = os.path.join(log_dir, f"traffic_stats_{timestamp}.csv")
        with open(csv_file, 'w', encoding='utf-8') as f:
            # 写入表头
            f.write("文件名,时间戳,汽车数量,处理时间(秒),警告信息\n")
            
            # 写入每行数据
            for result in stats['image_results']:
                warnings_str = " | ".join(result['warnings']) if result['warnings'] else "无"
                f.write(f"{result['filename']},{result['timestamp']},{result['car_count']},{result['processing_time']:.4f},{warnings_str}\n")
                
                # 收集Excel数据（只包含图片名、时间和车辆数）
                excel_data.append({
                    '图片名称': result['filename'],
                    '处理时间': result['timestamp'],
                    '车辆总数': result['car_count']
                })
        
        # 保存摘要统计到文本文件
        summary_file = os.path.join(log_dir, f"traffic_summary_{timestamp}.txt")
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write("车流量统计摘要\n")
            f.write(f"统计时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.write(f"处理图片总数: {stats['total_images']}\n")
            f.write(f"检测到汽车总数: {stats['total_cars']}\n")
            f.write(f"平均每图片汽车数: {stats['avg_cars_per_image']:.2f}\n")
            f.write(f"总处理时间: {stats['total_processing_time']:.2f} 秒\n")
            f.write(f"平均处理时间: {stats['avg_processing_time']:.2f} 秒/图片\n")
        
        # 保存到Excel文件
        excel_file = os.path.join(log_dir, f"traffic_stats_{timestamp}.xlsx")
        try:
            # 使用pandas创建DataFrame并保存为Excel
            df = pd.DataFrame(excel_data)
            df.to_excel(excel_file, index=False, sheet_name='车辆统计')
            print(f"Excel统计报告已保存到: {excel_file}")
        except ImportError:
            # 如果pandas导入失败，尝试使用openpyxl
            try:
                wb = Workbook()
                ws = wb.active
                ws.title = "车辆统计"
                
                # 写入表头
                ws.append(['图片名称', '处理时间', '车辆总数'])
                
                # 写入数据
                for data in excel_data:
                    ws.append([data['图片名称'], data['处理时间'], data['车辆总数']])
                
                wb.save(excel_file)
                print(f"Excel统计报告已保存到: {excel_file}")
            except Exception as e:
                print(f"保存Excel文件时出错: {str(e)}")
        
        print(f"统计报告已保存到: {csv_file}")
        print(f"统计摘要已保存到: {summary_file}")


if __name__ == "__main__":
    # 初始化检测器
    detector = HighwayVehicleDetector()
    
    # 批量处理模式 - 处理指定文件夹中的所有图片
    image_folder = "E:\pythonPJ\Project2\imgs\hight_image"
    print(f"开始批量处理文件夹: {image_folder}")
    stats = detector.batch_process_images(image_folder)
    
    # 单张图片处理模式（保留原有功能，可注释掉批量处理部分来使用）
    # try:
    #     # 执行检测
    #     image_path = "E:\\pythonPJ\\Projext2\\imgs\\hight_image\\img.png"
    #     car_count, warnings, _ = detector.detect_and_process(image_path)
    # 
    #     # 输出结果
    #     print("\n===== 检测结果 =====")
    #     print(f"汽车总数：{car_count} 辆")
    # 
    #     if warnings:
    #         print("\n===== 非汽车提醒 =====")
    #         for i, warn in enumerate(warnings, 1):
    #             print(f"{i}. {warn}")
    #     else:
    #         print("\n未检测到非汽车目标，一切正常。")
    # 
    # except Exception as e:
    #     print(f"检测出错：{str(e)}")